import json
import subprocess
import threading

from flask import send_file, Blueprint, request, Flask, jsonify
import os
import configparser
import zipfile
import tarfile
import rarfile
import re


app = Flask(__name__)
bp = Blueprint('main', __name__, url_prefix='/model/inference')

# 定义服务器地址和端口
server_host = '0.0.0.0'
server_port = 8090

# 定义成功和失败响应体
success_response = {'staus': 200, 'message': 'success', 'data': {}}
failed_response = {'staus': 500, 'message': 'error', 'data': {}}

# 定义根目录
model_root_path = '/data/model/'
script_root_path = '/data/'
output_root_path = '/data/result/'

# 定义模型字典
model_dict = {
    'LongWriter-glm4-9b': f'{model_root_path}LongWriter-glm4-9b/',
    'Qwen1.5-7B-Chat': f'{model_root_path}Qwen1.5-7B-Chat/',
    'Qwen2-7B-Instruct': f'{model_root_path}Qwen2-7B-Instruct/',
    'Qwen2-VL-7B-Instruct': f'{model_root_path}Qwen2-VL-7B-Instruct/',
    'model_16b': f'{model_root_path}model_16b/',
    'model_33b': f'{model_root_path}model_33b/',
    'model_7b': f'{model_root_path}model_7b/',
    'codegeex4-all-9b': f'{model_root_path}codegeex4-all-9b/',
}

# 定义推理脚本路径
script_dict = {
    'LongWriter-glm4-9b': f'{script_root_path}inference_vllm_auto_insta_all8.py',
    'Qwen1.5-7B-Chat': f'{script_root_path}inference_vllm_auto_insta_all11.py',
    'Qwen2-7B-Instruct': f'{script_root_path}inference_vllm_auto_insta_all10.py',
    'Qwen2-VL-7B-Instruct': f'{script_root_path}inference_vllm_auto_insta_all12.py',
    'model_16b': f'{script_root_path}inference_vllm_auto_insta_all2.py',
    'model_33b': f'{script_root_path}inference_vllm_auto_insta_all3.py',
    'model_7b': f'{script_root_path}inference_vllm_auto_insta_all1.py',
    'codegeex4-all-9b': f'{script_root_path}inference_vllm_auto_insta_all9.py',
}

# 定义允许的文件扩展名
ALLOWED_EXTENSIONS = {'json', 'jsonl', 'zip', 'rar', 'tar.gz'}


def allowed_file(filename):
    # 检查文件扩展名是否在允许的列表中
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def zip_json_files(output_path, zip_filename):
    with zipfile.ZipFile(zip_filename, 'w') as zipf:
        for root, _, files in os.walk(output_path):
            for file in files:
                if file.endswith('.json'):
                    file_path = os.path.join(root, file)
                    zipf.write(file_path, os.path.relpath(file_path, output_path))


def unzip_file(zip_filepath, extract_to):
    with zipfile.ZipFile(zip_filepath, 'r') as zip_ref:
        for zip_info in zip_ref.infolist():
            # 手动处理文件名的编码问题
            zip_info.filename = zip_info.filename.encode('cp437').decode('gbk')
            zip_ref.extract(zip_info, extract_to)


def untar_file(tar_filepath, extract_to):
    if not os.path.exists(extract_to):
        os.makedirs(extract_to)

    with tarfile.open(tar_filepath, 'r:gz') as tar_ref:
        for tar_info in tar_ref.getmembers():
            if tar_info.isfile():
                # 手动处理文件名的编码问题
                tar_info.name = tar_info.name.encode('utf-8').decode('utf-8')
                tar_ref.extract(tar_info, path=extract_to)
        print(f"Files extracted to: {extract_to}")


def unrar_file(rar_filepath, extract_to):
    if not os.path.exists(extract_to):
        os.makedirs(extract_to)

    with rarfile.RarFile(rar_filepath, 'r') as rar_ref:
        for rar_info in rar_ref.infolist():
            # 手动处理文件名的编码问题
            rar_info.filename = rar_info.filename.encode('utf-8').decode('utf-8')
            rar_ref.extract(rar_info, path=extract_to)
        print(f"Files extracted to: {extract_to}")


def extract_file(filepath, extract_to):
    if not os.path.exists(extract_to):
        os.makedirs(extract_to)

    if filepath.endswith('.zip'):
        unzip_file(filepath, extract_to)
        os.remove(filepath)
    elif filepath.endswith('.tar.gz') or filepath.endswith('.tgz'):
        untar_file(filepath, extract_to)
        os.remove(filepath)
    elif filepath.endswith('.rar'):
        unrar_file(filepath, extract_to)
        os.remove(filepath)
    else:
        print(f"file format not need to extract: {filepath}")


def run_subprocess(cmd_list):
    process = subprocess.run(cmd_list, capture_output=True,
                             text=True)
    print(process.stdout)
    print(process.stderr)


@bp.route('/start', methods=['POST'])
def start():
    # 定义输出目录列表
    ouput_path_list = []
    try:
        # 获取请求参数
        models = request.form.get("model")
        inputPath = request.form.get("inputPath")
        tensorParallelSize = request.form.get("tensorParallelSize")
        batchSize = request.form.get("batchSize")
        inputLength = request.form.get("inputLength")
        outputLength = request.form.get("outputLength")

        # 如果缺失必填参数 返回错误
        if not models:
            failed_response['message'] = 'model is required'
            return jsonify(message=failed_response), 400
        if not inputPath:
            failed_response['message'] = 'inputPath is required'
            return jsonify(message=failed_response), 400
        # 非必填项设置默认值
        if not tensorParallelSize:
            tensorParallelSize = 8
        if not batchSize:
            batchSize = 16
        if not inputLength:
            inputLength = 8192
        if not outputLength:
            outputLength = 4096

        # 检查请求中是否包含文件部分
        if 'file' not in request.files:
            failed_response['message'] = 'No selected file'
            return failed_response, 400

        file = request.files['file']

        # 如果用户没有选择文件，浏览器也会提交一个空的文件名
        if file.filename == '':
            failed_response['message'] = 'No selected file'
            return failed_response, 400
        # 判断文件类型
        if file:
            if not allowed_file(file.filename):
                failed_response['message'] = 'File type not allowed'
                return failed_response, 400

        for model in models.split(','):
            # 创建目录
            if not os.path.exists(inputPath.join(model)):
                os.makedirs(inputPath.join(model))
                # 处理压缩文件
                if re.search(r'\.(zip|rar|tar.gz)$', file.filename):
                    extract_file(file.filename, inputPath)
                else:
                    # 保存文件到指定路径
                    file_path = os.path.join(inputPath, model, file.filename)
                    file.save(file_path)
            # 根据模型名启动对应的脚本
            cmd_list = ['python3', script_dict[model], '--model', model_dict[model], '--tokenizer', model_dict[model],
                        '--tensor-parallel-size', 8, '--output-len', 4096, '--batch-size', 16, '--num-iters', 1,
                        '--data-file', inputPath]
            print(f"cmd:{cmd_list}")
            # 创建一个线程并启动它
            # thread = threading.Thread(target=run_subprocess, args=(cmd_list,))
            # thread.start()

            # 生成一个finish文件用于记录任务状态
            with open(f'{output_root_path}/{model}/finish.txt', 'w') as file:
                # 写入内容到文件
                file.write(json.dumps({'status': 'Running', 'error': None}))
                ouput_path_list.append(f'{output_root_path}/{model}')
    except Exception as e:
        # 返回异常信息
        failed_response['message'] = str(e)
        return jsonify(message=failed_response)

    success_response['message'] = 'Start inference task success!'
    success_response['data']['outputPath'] = ouput_path_list
    return jsonify(message=success_response)


@bp.route('/status', methods=['POST'])
def status():
    # 获取请求参数
    param = request.get_json()
    outputPath = param.get('ouputPath')
    model = param.get('model')
    # 如果缺失必填参数 返回错误
    if not outputPath:
        failed_response['message'] = 'ouputPath is required'
        return jsonify(message=failed_response), 400
    if not model:
        failed_response['message'] = 'model is required'
        return jsonify(message=failed_response), 400

    # 构造 finish.txt 文件的完整路径
    file_path = os.path.join(f'{outputPath}/{model}/', 'finish.txt')

    try:
        # 读取 finish.txt 文件中的内容
        with open(file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)

        # 返回读取到的 JSON 数据
        success_response['data'] = data
        return jsonify(message=success_response)
    except FileNotFoundError:
        failed_response['message'] = f"File not found: {file_path}"
        return jsonify(message=failed_response), 400
    except json.JSONDecodeError:
        return jsonify({"error": "Invalid JSON format in finish.txt"}), 400


@bp.route('/result', methods=['POST'])
def result():
    # 获取请求参数
    data = request.get_json()
    outputPath = data.get("outputPath")
    model = data.get("model")
    if not outputPath:
        failed_response['message'] = 'outputPath is required'
        return jsonify(message=failed_response), 400
    if not model:
        failed_response['message'] = 'model is required'
        return jsonify(message=failed_response), 400

    # 构造 finish.txt 文件的完整路径
    finish_file_path = os.path.join(f'{outputPath}/{model}/finish.txt')

    try:
        # 读取 finish.txt 文件中的内容
        with open(finish_file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)

        # 解析 status 字段
        status = data.get('status')
        if status == 'success':
            # 获取输出目录中的所有 .json 文件
            json_files = [os.path.join(root, file)
                          for root, _, files in os.walk(f'{outputPath}/{model}')
                          for file in files if file.endswith('.json')]
            if len(json_files) == 1:
                # 只有一个 .json 文件，直接返回该文件
                return send_file(json_files[0], as_attachment=True)
            else:
                # 多个 .json 文件，压缩成 zip 并返回
                zip_filename = os.path.join(f'{outputPath}/{model}', 'results.zip')
                zip_json_files(f'{outputPath}/{model}', zip_filename)
                return send_file(zip_filename, as_attachment=True)
        elif status == 'running':
            success_response['data'] = data
            return jsonify(message=success_response), 200
        elif status == 'failed':
            failed_response['message'] = 'Task failed!'
            failed_response['data'] = data
            return jsonify(message=failed_response), 400
        else:
            failed_response['message'] = 'Unknown status in finish.txt'
            return jsonify(message=failed_response), 400

    except FileNotFoundError:
        failed_response['message'] = f"File not found: {finish_file_path}"
        return jsonify(message=failed_response), 404
    except json.JSONDecodeError:
        failed_response['message'] = 'Invalid JSON format in finish.txt'
        return jsonify(message=failed_response), 400


app.register_blueprint(bp)

if __name__ == '__main__':
    # config = configparser.ConfigParser()
    # config.read('config.ini')
    # server_host = config['server']['host']
    # server_port = int(config['server']['port'])
    app.run(debug=True, host=server_host, port=server_port)
